{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4zzlF9Mt_DSG"
      },
      "source": [
        "##### Copyright 2018 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "2FhUiw93_ITA"
      },
      "outputs": [],
      "source": [
        "#@title Default title text\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hymcIul0IzfO"
      },
      "source": [
        "# Simple, Distributed, and Accelerated Probabilistic Programming\n",
        "\n",
        "This notebook is a companion webpage for the NIPS 2018 paper, [\"Simple, Distributed, and Accelerated Probabilistic Programming\"](https://arxiv.org/abs/1811.02091) (Tran et al., 2018). See the [README.md](https://github.com/google-research/google-research/blob/master/simple_probabilistic_programming/README.md) for details on how to interact with data, models, probabilistic inference, and more. It assumes the following dependencies:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "8QpNWIwK-065"
      },
      "outputs": [],
      "source": [
        "!pip install scipy==1.0.0\n",
        "!pip install tensor2tensor==1.9.0\n",
        "!pip install tensorflow==1.12.0rc2  # alternatively, tensorflow-gpu==1.12.0rc2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "gpiiAsjJwMec"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import six\n",
        "import tensorflow as tf\n",
        "\n",
        "from contextlib import contextmanager\n",
        "from scipy import stats\n",
        "from tensor2tensor.layers import common_attention\n",
        "from tensor2tensor.layers import common_image_attention as cia\n",
        "from tensor2tensor.models import image_transformer as imgtransformer\n",
        "from tensor2tensor.models import transformer\n",
        "tfb = tf.contrib.distributions.bijectors\n",
        "tfe = tf.contrib.eager"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "N1HJVwBr2AXf"
      },
      "source": [
        "This notebook also requires importing files in this Github repository:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "DyTjay7j1_a4"
      },
      "outputs": [],
      "source": [
        "import simple_probabilistic_programming as ed\n",
        "from simple_probabilistic_programming import no_u_turn_sampler"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "jqaoZSgG0VVk"
      },
      "source": [
        "Certain snippets require eager execution. This is run with the command below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "bTjt6ZWo0Zpg"
      },
      "outputs": [],
      "source": [
        "tf.enable_eager_execution()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "9i64LOk2UvUW"
      },
      "source": [
        "## Section 2. Random Variables Are All You Need\n",
        "\n",
        "__Figure 1__. Beta-Bernoulli program. In eager mode, `model()` generates a binary vector of $50$ elements. In graph mode, `model()` returns an op to be evaluated in a TensorFlow session."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1cu3Byp86CSc"
      },
      "outputs": [],
      "source": [
        "def model():\n",
        "  p = ed.Beta(1., 1., name=\"p\")\n",
        "  x = ed.Bernoulli(probs=p, \n",
        "                   sample_shape=50, \n",
        "                   name=\"x\")\n",
        "  return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 68
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 874,
          "status": "ok",
          "timestamp": 1534123893617,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "Un2nvb7JJAa7",
        "outputId": "c21767d4-d126-44ff-9ea1-02151eddeb43"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "RandomVariable(\"\n",
            "[0 0 0 0 0 0 1 0 0 0 0 0 1 0 0 0 1 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0\n",
            " 0 0 0 0 0 0 0 0 0 0 0 0 0]\", shape=(50,), dtype=int32, device=/job:localhost/replica:0/task:0/device:CPU:0)\n"
          ]
        }
      ],
      "source": [
        "x = model()\n",
        "print(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0nRCFX2sVHvf"
      },
      "source": [
        "__Figure 2__. Variational program (Ranganath et al., 2016), available in eager mode. Python control flow is applicable to generative processes: given a coin flip, the program generates from one of two neural nets. Their outputs can have differing shape (and structure)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "zIxFB_Ni5_hv"
      },
      "outputs": [],
      "source": [
        "def neural_net_negative(noise, inputs):\n",
        "  net = noise + inputs\n",
        "  net = tf.layers.dense(net, 512, activation=tf.nn.relu)\n",
        "  net = tf.layers.dense(net, 64, activation=None)\n",
        "  return net\n",
        "\n",
        "def neural_net_positive(noise, inputs):\n",
        "  del noise, inputs  # unused\n",
        "  return \"Hello. I'm a different output type.\"\n",
        "\n",
        "def variational(x):\n",
        "  eps = ed.Normal(0., 1., sample_shape=2)\n",
        "  if eps[0] \u003e 0:\n",
        "    return neural_net_positive(eps[1], x)\n",
        "  else:\n",
        "    return neural_net_negative(eps[1], x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 34
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 488,
          "status": "ok",
          "timestamp": 1534128522889,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "nr3KLSQRVMzh",
        "outputId": "6a8456a3-3e7b-440f-d46b-79707147c377"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(\u003ctype 'EagerTensor'\u003e, TensorShape([Dimension(4), Dimension(64), Dimension(64), Dimension(64)]))\n"
          ]
        }
      ],
      "source": [
        "if not tf.executing_eagerly():\n",
        "  raise ValueError(\"This code snippet requires eager execution.\")\n",
        "\n",
        "x = tf.random_normal([4, 64, 64, 3])  # batch of, e.g., 64x64x3 images\n",
        "z = variational(x)\n",
        "if isinstance(z, tf.Tensor):\n",
        "  print(type(z), z.shape)  # to avoid printing a huge Tensor\n",
        "else:\n",
        "  print(z)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "vO0FPBywVv6P"
      },
      "source": [
        "__Figure 3.__ Distributed autoregressive flows. The default length is 8, each with 4 independent flows (Papamakarios et al., 2017). Each flow transforms inputs via layers respecting autoregressive ordering. Flows are partitioned across a virtual topology of 4x4 cores (rectangles); each core computes 2 flows and is locally connected; a final core aggregates. The virtual topology aligns with the physical TPU topology: for 4x4 TPUs, it is exact; for 16x16 TPUs, it is duplicated for data parallelism."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "iU1FsfyeVvHy"
      },
      "outputs": [],
      "source": [
        "class SplitAutoregressiveFlow(tfb.Bijector):\n",
        "  def __init__(flow_size=[4]*8):\n",
        "    self.flows = []\n",
        "    for num_splits in flow_size:\n",
        "      flow = SplitAutoregressiveFlow(masked_network, num_splits)\n",
        "      self.flows.append(flow)\n",
        "    self.flows.append(SplitAutoregressiveFlow(masked_network, 1))\n",
        "    super(SplitAutoregressiveFlow, self).__init__()\n",
        "\n",
        "  def _forward(self, x):\n",
        "    for l, flow in enumerate(self.flows):\n",
        "      with tf.device(tf.contrib.tpu.core(l%4)):\n",
        "        x = flow.forward(x)\n",
        "    return x\n",
        "\n",
        "  def _inverse_and_log_det_jacobian(self, y):\n",
        "    ldj = 0.\n",
        "    for l, flow in enumerate(self.flows[::-1]):\n",
        "      with tf.device(tf.contrib.tpu.core(l%4)):\n",
        "        y, new_ldj = flow.inverse_and_log_det_jacobian(y)\n",
        "        ldj += new_ldj\n",
        "    return y, ldj\n",
        "\n",
        "\n",
        "class DistributedAutoregressiveFlow(tfb.Bijector):\n",
        "  def __init__(flow_size=[4]*8):\n",
        "    self.flows = []\n",
        "    for num_splits in flow_size:\n",
        "      flow = SplitAutoregressiveFlow(masked_network, num_splits)\n",
        "      self.flows.append(flow)\n",
        "    self.flows.append(SplitAutoregressiveFlow(masked_network, 1))\n",
        "    super(DistributedAutoregressiveFlow, self).__init__()\n",
        "    \n",
        "  def _forward(self, x):\n",
        "    for l, flow in enumerate(self.flows):\n",
        "      with tf.device(tf.contrib.tpu.core(l%4)):\n",
        "        x = flow.forward(x)\n",
        "    return x\n",
        "\n",
        "  def _inverse_and_log_det_jacobian(self, y):\n",
        "    ldj = 0.\n",
        "    for l, flow in enumerate(self.flows[::-1]):\n",
        "      with tf.device(tf.contrib.tpu.core(l%4)):\n",
        "        y, new_ldj = flow.inverse_and_log_det_jacobian(y)\n",
        "        ldj += new_ldj\n",
        "    return y, ldj"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "tOTRYencV3uL"
      },
      "source": [
        "__Figure 4.__\n",
        "Model-parallel VAE with TPUs, generating 16-bit audio from 8-bit latents. The prior and decoder split computation according to distributed autoregressive flows. The encoder may split computation according to `compressor`; we omit it for space."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "78yfrFUPrzBG"
      },
      "outputs": [],
      "source": [
        "def prior():\n",
        "  \"\"\"Uniform noise to 8-bit latent, [u1,...,u(T/2)] -\u003e [z1,...,z(T/2)]\"\"\"\n",
        "  dist = ed.Independent(ed.Uniform(low=tf.zeros([batch_size, T/2])))\n",
        "  return ed.TransformedDistribution(dist, DistributedAutoregressiveFlow(flow_size))\n",
        "\n",
        "def decoder(z):\n",
        "  \"\"\"Uniform noise + latent to 16-bit audio, [u1,...,uT], [z1,...,z(T/2)] -\u003e [x1,...,xT]\"\"\"\n",
        "  dist = ed.Independent(ed.Uniform(low=tf.zeros([batch_size, T])))\n",
        "  dist = ed.TransformedDistribution(dist, tfb.Affine(shift=decompressor(z)))\n",
        "  return ed.TransformedDistribution(dist, DistributedAutoregressiveFlow(flow_size))\n",
        "\n",
        "def encoder(x):\n",
        "  \"\"\"16-bit audio to 8-bit latent, [x1,...,xT] -\u003e [z1,...,z(T/2)]\"\"\"\n",
        "  loc, log_scale = tf.split(compressor(x), 2, axis=-1)\n",
        "  return ed.MultivariateNormalDiag(loc=loc, scale=tf.exp(log_scale))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "RydhT_oasF-U"
      },
      "source": [
        "__Figure 5__. Edward2's core.\n",
        "`trace` defines a context; any traceable ops executed during it are replaced by calls to `tracer`. `traceable` registers these ops; we register Edward random variables."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "JZCaqZyLsPq-"
      },
      "outputs": [],
      "source": [
        "STACK = [lambda f, *a, **k: f(*a, **k)]\n",
        "\n",
        "@contextmanager\n",
        "def trace(tracer):\n",
        "  STACK.append(tracer)\n",
        "  yield\n",
        "  STACK.pop()\n",
        "  \n",
        "def traceable(f):\n",
        "  def f_wrapped(*a, **k):\n",
        "    STACK[-1](f, *a, **k)\n",
        "  return f_wrapped"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "YkxhosqpsSp-"
      },
      "source": [
        "__Figure 7__. A higher-order function which takes a `model` program as input and returns its log-joint density function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "_1PHF1XGn1l3"
      },
      "outputs": [],
      "source": [
        "def make_log_joint_fn(model):\n",
        "  def log_joint_fn(**model_kwargs):\n",
        "    def tracer(rv_call, *args, **kwargs):\n",
        "      name = kwargs.get(\"name\")\n",
        "      kwargs[\"value\"] = model_kwargs.get(name)\n",
        "      rv = rv_call(*args, **kwargs)\n",
        "      log_probs.append(tf.reduce_sum(rv.distribution.log_prob(rv)))\n",
        "      return rv\n",
        "    log_probs = []\n",
        "    with ed.trace(tracer):\n",
        "      model()\n",
        "    return sum(log_probs)\n",
        "  return log_joint_fn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 34
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 929,
          "status": "ok",
          "timestamp": 1534123949958,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "kIiCW84zsmrT",
        "outputId": "f22a3da2-0731-4285-cc96-47db3471d898"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tf.Tensor(-23.1994, shape=(), dtype=float32)\n"
          ]
        }
      ],
      "source": [
        "try:\n",
        "  model\n",
        "except NameError:\n",
        "  raise NameError(\"This code snippet requires `model` from above.\")\n",
        "\n",
        "log_joint = make_log_joint_fn(model)\n",
        "p = np.random.uniform()\n",
        "x = np.round(np.random.normal(size=[50])).astype(np.int32)\n",
        "out = log_joint(p=p, x=x)\n",
        "print(out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "KHaHv3OosqkA"
      },
      "source": [
        "__Figure 8__. A higher-order function which takes a `model` program as input and returns its causally intervened program. Intervention differs from conditioning: it does not change the sampled value but the distribution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "yUNzFmeen_1h"
      },
      "outputs": [],
      "source": [
        "def mutilate(model, **do_kwargs):\n",
        "  def mutilated_model(*args, **kwargs):\n",
        "    def tracer(rv_call, *args, **kwargs):\n",
        "      name = kwargs.get(\"name\")\n",
        "      if name in do_kwargs:\n",
        "        return do_kwargs[name]\n",
        "      return rv_call(*args, **kwargs)\n",
        "    with ed.trace(tracer):\n",
        "      return model(*args, **kwargs)\n",
        "  return mutilated_model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 68
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 53,
          "status": "ok",
          "timestamp": 1534122871529,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "2b2HfR6ks6EM",
        "outputId": "9c5df1e3-e792-4081-a0ef-05f804631f32"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "RandomVariable(\"\n",
            "[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1\n",
            " 1 1 1 1 1 1 1 1 1 1 1 1 1]\", shape=(50,), dtype=int32, device=/job:localhost/replica:0/task:0/device:CPU:0)\n"
          ]
        }
      ],
      "source": [
        "try:\n",
        "  model\n",
        "except NameError:\n",
        "  raise NameError(\"This code snippet requires `model` from above.\")\n",
        "\n",
        "mutilated_model = mutilate(model, p=0.999)\n",
        "x = mutilated_model()\n",
        "print(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8tJY1LVRs8d4"
      },
      "source": [
        "## Section 3. Learning with Low-Level Functions\n",
        "\n",
        "__FIgure 9__. Data-parallel Image Transformer with TPUs (Parmar et al., 2018). It is a neural autoregressive model which computes the log-probability of a batch of images with self-attention. Edward2 enables representing and training the model as a log-probability function; this is more efficient than the typical representation of programs as a generative process."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "dyOZR0DFoHKA"
      },
      "outputs": [],
      "source": [
        "get_channel_embeddings = cia.get_channel_embeddings\n",
        "add_positional_embedding = common_attention.add_positional_embedding\n",
        "local_attention_1d = cia.local_attention_1d\n",
        "\n",
        "def image_transformer(inputs, hparams):\n",
        "  x = get_channel_embeddings(3, inputs, hparams.hidden_size)\n",
        "  x = tf.reshape(x, [-1, 32*32*3, hparams.hidden_size])\n",
        "  x = tf.pad(x, [[0, 0], [1, 0], [0, 0]])[:, :-1, :]  # shift pixels right\n",
        "  x = add_positional_embedding(x, max_length=32*32*3+3, name=\"pos_embed\")\n",
        "  x = tf.nn.dropout(x, keep_prob=0.7)\n",
        "  for _ in range(hparams.num_decoder_layers):\n",
        "    with tf.variable_scope(None, default_name=\"decoder_layer\"):\n",
        "      y = local_attention_1d(x, hparams, attention_type=\"local_mask_right\", \n",
        "                             q_padding=\"LEFT\", kv_padding=\"LEFT\")\n",
        "      x = tf.contrib.layers.layer_norm(tf.nn.dropout(y, keep_prob=0.7) + x, begin_norm_axis=-1)\n",
        "      y = tf.layers.dense(x, hparams.filter_size, activation=tf.nn.relu)\n",
        "      y = tf.layers.dense(y, hparams.hidden_size, activation=None)\n",
        "      x = tf.contrib.layers.layer_norm(tf.nn.dropout(y, keep_prob=0.7) + x, begin_norm_axis=-1)\n",
        "  x = tf.reshape(x, [-1, 32, 32, 3, hparams.hidden_size])\n",
        "  logits = tf.layers.dense(x, 256, activation=None)\n",
        "  return ed.Categorical(logits=logits).distribution.log_prob(inputs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 35
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 12084,
          "status": "ok",
          "timestamp": 1534177612610,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "0m9fZ5APtLaP",
        "outputId": "977a682d-3de8-4fe1-b4d9-26e4e09f9b69"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Tensor(\"Neg:0\", shape=(), dtype=float32)\n"
          ]
        }
      ],
      "source": [
        "if tf.executing_eagerly():\n",
        "  raise ValueError(\"This code snippet does not support eager execution.\")\n",
        "\n",
        "batch_size = 4\n",
        "inputs = tf.random_uniform([batch_size, 32, 32, 3], minval=0, maxval=256, dtype=tf.int32)\n",
        "hparams = imgtransformer.imagetransformer_cifar10_base()\n",
        "loss = -tf.reduce_sum(image_transformer(inputs, hparams))\n",
        "train_op = tf.contrib.tpu.CrossShardOptimizer(tf.train.AdamOptimizer()).minimize(loss)\n",
        "print(loss)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-aUoyMGvho2l"
      },
      "source": [
        "__Figure 10__. Core logic in No-U-Turn Sampler (Hoffman and Gelman, 2014).  This algorithm has data-dependent non-tail recursion.\n",
        "\n",
        "See [`no_u_turn_sampler/`](https://github.com/google-research/google-research/tree/master/simple_probabilistic_programming/no_u_turn_sampler/) in the Github repository for its full implementation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "scNQtYbTtXlQ"
      },
      "source": [
        "__Figure 11__. Variational inference with preconditioned gradient descent. Edward2 offers writing the probabilistic program and performing arbitrary TensorFlow computation for learning."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "kty4nu5c1fQ1"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  model\n",
        "  make_log_joint_fn\n",
        "except NameError:\n",
        "  raise NameError(\"This code snippet requires `model`, `make_log_joint_fn` \"\n",
        "                  \" from above.\")\n",
        "\n",
        "class Variational(object):\n",
        "  def __init__(self):\n",
        "    self.parameters = tf.random_normal([2])\n",
        "  \n",
        "  def __call__(self, x):\n",
        "    del x  # unused; it is a non-amortized approximation\n",
        "    return ed.Deterministic(loc=tf.sigmoid(self.parameters[0]), name=\"qp\")\n",
        "\n",
        "variational = Variational()\n",
        "x = tf.random_uniform([50], minval=0, maxval=2, dtype=tf.int32)\n",
        "alignment = {\"qp\": \"p\"}\n",
        "\n",
        "def loss(x):\n",
        "  qz = variational(x)\n",
        "  log_joint_fn = make_log_joint_fn(model)\n",
        "  kwargs = {alignment[rv.distribution.name]: rv.value\n",
        "            for rv in [qz]}\n",
        "  energy = log_joint_fn(x=x, **kwargs)\n",
        "  entropy = sum([rv.distribution.entropy() for rv in [qz]])\n",
        "  return -energy - entropy\n",
        "\n",
        "def grad():\n",
        "  with tf.GradientTape() as tape:\n",
        "    tape.watch(variational.parameters)\n",
        "    loss_value = loss(x)\n",
        "  return tape.gradient(loss_value, variational.parameters)\n",
        "\n",
        "def train(precond):\n",
        "  for _ in range(5):\n",
        "    grads = tf.tensordot(precond, grad(), [[1], [0]])\n",
        "    variational.parameters -= 0.1 * grads\n",
        "  return loss(x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 34
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 189,
          "status": "ok",
          "timestamp": 1534128427976,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "galdOECztcXE",
        "outputId": "5c822295-d9c2-4bc9-bb71-2c26e7eb4184"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tf.Tensor(34.2965, shape=(), dtype=float32)\n"
          ]
        }
      ],
      "source": [
        "if not tf.executing_eagerly():\n",
        "  raise ValueError(\"This code snippet requires eager execution.\")\n",
        "\n",
        "precond = tf.eye(2)\n",
        "loss_value = train(precond)\n",
        "print(loss_value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "pjWpGC22tfJK"
      },
      "source": [
        "__Figure 12__. Learning-to-learn. It finds the optimal preconditioner for `train` (__Figure 11__) by differentiating the entire learning algorithm with respect to the preconditioner."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 357
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 1240,
          "status": "ok",
          "timestamp": 1534128429769,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "c8Or6wtdtl0F",
        "outputId": "ff61f5e6-202f-4cd5-c681-ec8787cdc981"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(34.296486, array([[ 0.67495859,  0.81464583],\n",
            "       [-0.18986063, -0.0458361 ]], dtype=float32))\n",
            "(34.296494, array([[ 0.04607767,  0.81464583],\n",
            "       [-0.18986063, -0.0458361 ]], dtype=float32))\n",
            "(34.296486, array([[-0.44003215,  0.81464583],\n",
            "       [-0.18986063, -0.0458361 ]], dtype=float32))\n",
            "(34.296486, array([[-0.83821177,  0.81464583],\n",
            "       [-0.18986063, -0.0458361 ]], dtype=float32))\n",
            "(34.296486, array([[-1.1741091 ,  0.81464583],\n",
            "       [-0.18986063, -0.0458361 ]], dtype=float32))\n",
            "(34.296501, array([[-0.69452661,  0.81464583],\n",
            "       [-0.18986063, -0.0458361 ]], dtype=float32))\n",
            "(34.300507, array([[-0.18749332,  0.81464583],\n",
            "       [-0.18986063, -0.0458361 ]], dtype=float32))\n",
            "(34.328743, array([[ 0.34323317,  0.81464583],\n",
            "       [-0.18986063, -0.0458361 ]], dtype=float32))\n",
            "(34.296623, array([[ 0.81854713,  0.81464583],\n",
            "       [-0.18986063, -0.0458361 ]], dtype=float32))\n",
            "(34.296486, array([[ 1.24275446,  0.81464583],\n",
            "       [-0.18986063, -0.0458361 ]], dtype=float32))\n"
          ]
        }
      ],
      "source": [
        "if not tf.executing_eagerly():\n",
        "  raise ValueError(\"This code snippet requires eager execution.\")\n",
        "  \n",
        "precond = tfe.Variable(tf.random_normal([2, 2]))\n",
        "optimizer = tf.train.AdamOptimizer(1.)\n",
        "for _ in range(10):\n",
        "  with tf.GradientTape() as tape:\n",
        "    loss_value = train(precond)\n",
        "  grads = tape.gradient(loss_value, [precond])\n",
        "  optimizer.apply_gradients(zip(grads, [precond]))\n",
        "  print(loss_value.numpy(), precond.numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ZREd7iIEtnWO"
      },
      "source": [
        "## Appendix A. Edward2 on SciPy\n",
        "\n",
        "We  illustrate  the  broad  applicability  of  Edward2’s  tracing  by  implementing  Edward2  on  top  of SciPy.\n",
        "\n",
        "For this notebook, we mimick a namespace using a struct so that one can play with the traceable scipy stats here."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 55
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 33,
          "status": "ok",
          "timestamp": 1534180922861,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "VHgHRftCtuId",
        "outputId": "1f44c650-6de1-4607-b303-483f8e3005f1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "['alpha', 'anglit', 'arcsine', 'argus', 'bernoulli', 'beta', 'betaprime', 'binom', 'boltzmann', 'bradford', 'burr', 'burr12', 'cauchy', 'chi', 'chi2', 'cosine', 'crystalball', 'dgamma', 'dirichlet', 'dlaplace', 'dweibull', 'erlang', 'expon', 'exponnorm', 'exponpow', 'exponweib', 'f', 'fatiguelife', 'fisk', 'foldcauchy', 'foldnorm', 'frechet_l', 'frechet_r', 'gamma', 'gausshyper', 'genexpon', 'genextreme', 'gengamma', 'genhalflogistic', 'genlogistic', 'gennorm', 'genpareto', 'geom', 'gilbrat', 'gompertz', 'gumbel_l', 'gumbel_r', 'halfcauchy', 'halfgennorm', 'halflogistic', 'halfnorm', 'hypergeom', 'hypsecant', 'invgamma', 'invgauss', 'invweibull', 'invwishart', 'johnsonsb', 'johnsonsu', 'kappa3', 'kappa4', 'ksone', 'kstwobign', 'laplace', 'levy', 'levy_l', 'levy_stable', 'loggamma', 'logistic', 'loglaplace', 'lognorm', 'logser', 'lomax', 'matrix_normal', 'maxwell', 'mielke', 'multinomial', 'multivariate_normal', 'nakagami', 'nbinom', 'ncf', 'nct', 'ncx2', 'norm', 'ortho_group', 'pareto', 'pearson3', 'planck', 'poisson', 'powerlaw', 'powerlognorm', 'powernorm', 'randint', 'random_correlation', 'rayleigh', 'rdist', 'recipinvgauss', 'reciprocal', 'rice', 'semicircular', 'skellam', 'skewnorm', 'special_ortho_group', 't', 'trapz', 'triang', 'truncexpon', 'truncnorm', 'tukeylambda', 'uniform', 'unitary_group', 'vonmises', 'vonmises_line', 'wald', 'weibull_max', 'weibull_min', 'wishart', 'wrapcauchy', 'zipf']\n"
          ]
        }
      ],
      "source": [
        "class FakeEdward2ScipyNamespace(object):\n",
        "  pass\n",
        "\n",
        "for _name in sorted(dir(stats)):\n",
        "  _candidate = getattr(stats, _name)\n",
        "  if isinstance(_candidate, (stats._multivariate.multi_rv_generic,\n",
        "                             stats.rv_continuous,\n",
        "                             stats.rv_discrete,\n",
        "                             stats.rv_histogram)):\n",
        "    _candidate.rvs = ed.traceable(_candidate.rvs)\n",
        "    setattr(FakeEdward2ScipyNamespace, _name, _candidate)\n",
        "    del _candidate\n",
        "        \n",
        "scipy_stats = FakeEdward2ScipyNamespace()\n",
        "print([name for name in dir(scipy_stats) if not name.startswith(\"__\")])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "FVQQih5cttVB"
      },
      "source": [
        "Below is an Edward2 linear regression program on SciPy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "k2S_NCFxoVpN"
      },
      "outputs": [],
      "source": [
        "def make_log_joint_fn(model):\n",
        "  def log_joint_fn(*model_args, **model_kwargs):\n",
        "    def tracer(rv_call, *args, **kwargs):\n",
        "      name = kwargs.pop(\"name\", None)\n",
        "      kwargs.pop(\"size\", None)\n",
        "      kwargs.pop(\"random_state\", None)\n",
        "      value = model_kwargs.get(name)\n",
        "      log_prob_fn = getattr(scipy_stats, rv_call.im_class.__name__[:-4]).logpdf\n",
        "      log_prob = np.sum(log_prob_fn(value, *args, **kwargs))\n",
        "      log_probs.append(log_prob)\n",
        "      return value\n",
        "    log_probs = []\n",
        "    with ed.trace(tracer):\n",
        "      model(*model_args)\n",
        "    return sum(log_probs)\n",
        "  return log_joint_fn\n",
        "\n",
        "def linear_regression(X):\n",
        "  beta = scipy_stats.norm.rvs(loc=0.0, scale=0.1, size=X.shape[1], name=\"beta\")\n",
        "  loc = np.einsum('ij,j-\u003ei', X, beta)\n",
        "  y = scipy_stats.norm.rvs(loc=loc, scale=1., size=1, name=\"y\")\n",
        "  return y"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 35
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 48,
          "status": "ok",
          "timestamp": 1534181166389,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "PIXJPaRUt2Eo",
        "outputId": "b1b34c9e-0d92-4584-a9ea-1100b4c62e87"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "-1.0854508813\n"
          ]
        }
      ],
      "source": [
        "log_joint = make_log_joint_fn(linear_regression)\n",
        "\n",
        "X = np.random.normal(size=[3, 2])\n",
        "beta = np.random.normal(size=[2])\n",
        "y = np.random.normal(size=[3])\n",
        "out = log_joint(X, beta=beta, y=y)\n",
        "print(out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DMEm6RCSt335"
      },
      "source": [
        "## Appendix B. Grammar Variational Auto-Encoder\n",
        "\n",
        "The grammar variational auto-encoder (VAE) (Kusner et al., 2017) posits a generative model over\n",
        "productions from a context-free grammar, and it posits an amortized variational\n",
        "approximation for efficient posterior inference. We train the grammar VAE\n",
        "on synthetic data using the grammar from Kusner et al. (2017; Figure 1).\n",
        "\n",
        "This example showcases eager execution in order to train the model where data\n",
        "points have a variable number of time steps. However, note that this requires a\n",
        "batch size of 1. In this example, we assume data points arrive in a stream, one\n",
        "at a time. Such a setting requires handling a variable number of time steps\n",
        "as the maximum length is unbounded."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "nWppHx2rxrUy"
      },
      "outputs": [],
      "source": [
        "class SmilesGrammar(object):\n",
        "  \"\"\"Context-free grammar for SMILES strings.\"\"\"\n",
        "  nonterminal_symbols = {\"smiles\", \"chain\", \"branched atom\", \"atom\", \"ringbond\",\n",
        "                         \"aromatic organic\", \"aliphatic organic\", \"digit\"}\n",
        "  alphabet = {\"c\", \"C\", \"N\", \"1\", \"2\"}\n",
        "  production_rules = [\n",
        "      (\"smiles\", [\"chain\"]),\n",
        "      (\"chain\", [\"chain\", \"branched atom\"]),\n",
        "      (\"chain\", [\"branched atom\"]),\n",
        "      (\"branched atom\", [\"atom\", \"ringbond\"]),\n",
        "      (\"branched atom\", [\"atom\"]),\n",
        "      (\"atom\", [\"aromatic organic\"]),\n",
        "      (\"atom\", [\"aliphatic organic\"]),\n",
        "      (\"ringbond\", [\"digit\"]),\n",
        "      (\"aromatic organic\", [\"c\"]),\n",
        "      (\"aliphatic organic\", [\"C\"]),\n",
        "      (\"aliphatic organic\", [\"N\"]),\n",
        "      (\"digit\", [\"1\"]),\n",
        "      (\"digit\", [\"2\"]),\n",
        "  ]\n",
        "  start_symbol = \"smiles\"\n",
        "\n",
        "  def mask(self, symbol, on_value=0., off_value=-1e9):\n",
        "    \"\"\"Produces a masking tensor for (in)valid production rules.\"\"\"\n",
        "    mask_values = []\n",
        "    for lhs, _ in self.production_rules:\n",
        "      if symbol in lhs:\n",
        "        mask_value = on_value\n",
        "      else:\n",
        "        mask_value = off_value\n",
        "      mask_values.append(mask_value)\n",
        "    mask_values = tf.reshape(mask_values, [1, len(self.production_rules)])\n",
        "    return mask_values\n",
        "\n",
        "class ProbabilisticGrammar(tf.keras.Model):\n",
        "  \"\"\"Deep generative model over productions which follow a grammar.\"\"\"\n",
        "\n",
        "  def __init__(self, grammar, latent_size, num_units):\n",
        "    \"\"\"Constructs a probabilistic grammar.\"\"\"\n",
        "    super(ProbabilisticGrammar, self).__init__()\n",
        "    self.grammar = grammar\n",
        "    self.latent_size = latent_size\n",
        "    self.lstm = tf.nn.rnn_cell.LSTMCell(num_units)\n",
        "    self.output_layer = tf.keras.layers.Dense(len(grammar.production_rules))\n",
        "\n",
        "  def call(self, inputs):\n",
        "    \"\"\"Runs the model forward to generate a sequence of productions.\"\"\"\n",
        "    del inputs  # unused\n",
        "    latent_code = ed.MultivariateNormalDiag(loc=tf.zeros(self.latent_size),\n",
        "                                            sample_shape=1,\n",
        "                                            name=\"latent_code\")\n",
        "    state = self.lstm.zero_state(1, dtype=tf.float32)\n",
        "    t = 0\n",
        "    productions = []\n",
        "    stack = [self.grammar.start_symbol]\n",
        "    while stack:\n",
        "      symbol = stack.pop()\n",
        "      net, state = self.lstm(latent_code, state)\n",
        "      logits = self.output_layer(net) + self.grammar.mask(symbol)\n",
        "      production = ed.OneHotCategorical(logits=logits,\n",
        "                                        name=\"production_\" + str(t))\n",
        "      _, rhs = self.grammar.production_rules[tf.argmax(production, axis=1)]\n",
        "      for symbol in rhs:\n",
        "        if symbol in self.grammar.nonterminal_symbols:\n",
        "          stack.append(symbol)\n",
        "      productions.append(production)\n",
        "      t += 1\n",
        "    return tf.stack(productions, axis=1)\n",
        "\n",
        "class ProbabilisticGrammarVariational(tf.keras.Model):\n",
        "  \"\"\"Amortized variational posterior for a probabilistic grammar.\"\"\"\n",
        "\n",
        "  def __init__(self, latent_size):\n",
        "    \"\"\"Constructs a variational posterior for a probabilistic grammar.\"\"\"\n",
        "    super(ProbabilisticGrammarVariational, self).__init__()\n",
        "    self.latent_size = latent_size\n",
        "    self.encoder_net = tf.keras.Sequential([\n",
        "        tf.keras.layers.Conv1D(64, 3, padding=\"SAME\"),\n",
        "        tf.keras.layers.BatchNormalization(),\n",
        "        tf.keras.layers.Activation(tf.nn.elu),\n",
        "        tf.keras.layers.Conv1D(128, 3, padding=\"SAME\"),\n",
        "        tf.keras.layers.BatchNormalization(),\n",
        "        tf.keras.layers.Activation(tf.nn.elu),\n",
        "        tf.keras.layers.Dropout(0.1),\n",
        "        tf.keras.layers.GlobalAveragePooling1D(),\n",
        "        tf.keras.layers.Dense(latent_size * 2, activation=None),\n",
        "    ])\n",
        "\n",
        "  def call(self, inputs):\n",
        "    \"\"\"Runs the model forward to return a stochastic encoding.\"\"\"\n",
        "    net = self.encoder_net(tf.cast(inputs, tf.float32))\n",
        "    return ed.MultivariateNormalDiag(\n",
        "        loc=net[..., :self.latent_size],\n",
        "        scale_diag=tf.nn.softplus(net[..., self.latent_size:]),\n",
        "        name=\"latent_code_posterior\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 377
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 463,
          "status": "ok",
          "timestamp": 1534270019355,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "T_t1_SzAxq_y",
        "outputId": "d902e898-5776-434f-c08f-2dad813b98ff"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Production Shape: (1, 5, 13)\n",
            "String: N\n",
            "Encoded Productions: [[-0.20581727  0.92976463 -0.38938859  0.19347586 -0.74668086 -0.39914176\n",
            "   0.93074167  0.36102515]]\n",
            "Production Shape: (1, 8, 13)\n",
            "String: 2N\n",
            "Encoded Productions: [[-1.10246503  0.50674993  0.97561336 -1.5284797  -0.20616516  0.84167266\n",
            "   0.78415519 -0.26026636]]\n",
            "Production Shape: (1, 13, 13)\n",
            "String: 2N2C\n",
            "Encoded Productions: [[ 0.67856091  1.32186925 -0.84166789  0.60194713 -0.94573712 -0.85340422\n",
            "   0.36050722 -0.9566167 ]]\n",
            "Production Shape: (1, 24, 13)\n",
            "String: 1N21c2N\n",
            "Encoded Productions: [[-0.80259132  0.27463835 -0.45711571  0.42127368 -0.3014937  -0.09925826\n",
            "   0.34425485  0.13968848]]\n",
            "Production Shape: (1, 13, 13)\n",
            "String: 221N\n",
            "Encoded Productions: [[ 0.06864282  0.02749904  0.06260949  0.24946308  0.1729231   0.04198643\n",
            "   0.26645684  0.8167249 ]]\n"
          ]
        }
      ],
      "source": [
        "if not tf.executing_eagerly():\n",
        "  raise ValueError(\"This code snippet requires eager execution.\")\n",
        "\n",
        "grammar = SmilesGrammar()\n",
        "probabilistic_grammar = ProbabilisticGrammar(\n",
        "    grammar=grammar, latent_size=8, num_units=128)\n",
        "probabilistic_grammar_variational = ProbabilisticGrammarVariational(\n",
        "    latent_size=8)\n",
        "\n",
        "for _ in range(5):\n",
        "  productions = probabilistic_grammar(_)\n",
        "  print(\"Production Shape: {}\".format(productions.shape))\n",
        "\n",
        "  string = grammar.convert_to_string(productions)\n",
        "  print(\"String: {}\".format(string))\n",
        "\n",
        "  encoded_production = probabilistic_grammar_variational(productions)\n",
        "  print(\"Encoded Productions: {}\".format(encoded_production.numpy()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "jl-c7XL8xquQ"
      },
      "source": [
        "See [`tensorflow_probability/examples/grammar_vae.py`](https://github.com/tensorflow/probability/blob/master/tensorflow_probability/examples/grammar_vae.py) for the full example."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "swsn3og7uCZQ"
      },
      "source": [
        "## Appendix C. Markov chain Monte Carlo within Variational Inference\n",
        "\n",
        "\n",
        "We demonstrate another level of composability: inference within a probabilistic program. Namely, we apply MCMC to construct a flexible family of distributions for variational inference\n",
        "(Salimans et al., 2015; Hoffman, 2017). We apply a chain of transition kernels specified by NUTS (`nuts`) in Section 3.3 and the variational inference algorithm specified by `train` in __Figure 12__."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "qCMF_X-LuJmE"
      },
      "outputs": [],
      "source": [
        "class DeepLatentGaussianModel(tf.keras.Model):\n",
        "  \"\"\"Deep generative model.\"\"\"\n",
        "  def __init__(self, latent_size, data_shape, batch_size):\n",
        "    super(DeepLatentGaussianModel, self).__init__()\n",
        "    self.latent_size = latent_size\n",
        "    self.data_shape = data_shape\n",
        "    self.batch_size = batch_size\n",
        "    self.decoder_net = tf.keras.Sequential([\n",
        "        tf.keras.layers.Dense(512, activation=tf.nn.relu),\n",
        "        tf.keras.layers.Dense(np.prod(data_shape), activation=None),\n",
        "        tf.keras.layers.Reshape(data_shape),\n",
        "    ])\n",
        "\n",
        "  def call(self, inputs):\n",
        "    del inputs  # unused\n",
        "    latent_code = ed.MultivariateNormalDiag(\n",
        "        loc=tf.zeros([self.batch_size, self.latent_size]),\n",
        "        scale_diag=tf.ones([self.batch_size, self.latent_size]),\n",
        "        name=\"latent_code\")\n",
        "    data = ed.Categorical(logits=self.decoder_net(latent_code), name=\"data\")\n",
        "    return data\n",
        "\n",
        "class DeepLatentGaussianModelVariational(tf.keras.Model):\n",
        "  \"\"\"Amortized variational posterior.\"\"\"\n",
        "  def __init__(self,\n",
        "               latent_size,\n",
        "               data_shape,\n",
        "               num_transitions,\n",
        "               target_log_prob_fn,\n",
        "               step_size):\n",
        "    super(DeepLatentGaussianModelVariational, self).__init__()\n",
        "    self.latent_size = latent_size\n",
        "    self.data_shape = data_shape\n",
        "    self.num_transitions = num_transitions\n",
        "    self.target_log_prob_fn = target_log_prob_fn\n",
        "    self.step_size = step_size\n",
        "    self.encoder_net = tf.keras.Sequential([\n",
        "        tf.keras.layers.Reshape(np.prod(data_shape)),\n",
        "        tf.keras.layers.Dense(512, activation=tf.nn.relu),\n",
        "        tf.keras.layers.Dense(latent_size * 2, activation=None),\n",
        "    ])\n",
        "  \n",
        "  def call(self, inputs):\n",
        "    net = encoder_net(inputs)\n",
        "    qz = ed.MultivariateNormalDiag(\n",
        "        loc=net[..., :self.latent_size],\n",
        "        scale_diag=tf.nn.softplus(net[..., self.latent_size:]),\n",
        "        name=\"latent_code_posterior\")\n",
        "    current_target_log_prob = None\n",
        "    current_grads_target_log_prob = None\n",
        "    for _ in range(self.num_transitions):\n",
        "      [\n",
        "          [qz],\n",
        "          current_target_log_prob,\n",
        "          current_grads_target_log_prob,\n",
        "      ] = self._kernel(\n",
        "          current_state=[qz],\n",
        "          current_target_log_prob=current_target_log_prob,\n",
        "          current_grads_target_log_prob=current_grads_target_log_prob)\n",
        "    return qz\n",
        " \n",
        "  def _kernel(self, current_state, current_target_log_prob,\n",
        "              current_grads_target_log_prob):\n",
        "    return no_u_turn_sampler.kernel(\n",
        "        current_state=current_state,\n",
        "        target_log_prob_fn=self.target_log_prob_fn,\n",
        "        step_size=self.step_size,\n",
        "        current_target_log_prob=current_target_log_prob,\n",
        "        current_grads_target_log_prob=current_grads_target_log_prob)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "N9dIuksNwSa8"
      },
      "outputs": [],
      "source": [
        "latent_size = 50\n",
        "data_shape = [32, 32, 3, 256]\n",
        "batch_size = 4\n",
        "\n",
        "features = tf.random_normal([batch_size] + data_shape)\n",
        "model = DeepLatentGaussianModel(\n",
        "    latent_size=latent_size,\n",
        "    data_shape=data_shape,\n",
        "    batch_size=batch_size)\n",
        "variational = DeepLatentGaussianModelVariational(\n",
        "    latent_size=latent_size,\n",
        "    data_shape=data_shape,\n",
        "    step_size=[0.1],\n",
        "    target_log_prob_fn=lambda z: ed.make_log_joint(model)(x=x, z=z)\n",
        "    num_transitions=10)\n",
        "alignment = {\"latent_code_posterior\": \"latent_code\"}\n",
        "optimizer = tf.train.AdamOptimizer(1e-2)\n",
        "\n",
        "for _ in range(10):\n",
        "  with tf.GradientTape() as tape:\n",
        "    with ed.trace() as variational_tape:\n",
        "      _ = variational(features)\n",
        "    log_joint_fn = make_log_joint_fn(model)\n",
        "    kwargs = {alignment[rv.distribution.name]: rv.value\n",
        "              for rv in six.itervalues(variational_tape)}\n",
        "    energy = log_joint_fn(data=features, **kwargs)\n",
        "    entropy = sum([rv.distribution.entropy()\n",
        "                   for rv in six.itervalues(variational_tape)])\n",
        "    loss_value = -energy - entropy\n",
        "  grads = tape.gradient(loss_value, variational.variables)\n",
        "  optimizer.apply_gradients(zip(grads, variational.variables))\n",
        "  print(\"Step: {:\u003e3d} Loss: {:.3f}\".format(step, loss_value))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "PVZecP4fuKZZ"
      },
      "source": [
        "## Appendix D. No-U-Turn Sampler\n",
        "\n",
        "We implement an Edward program for Bayesian logistic regression with NUTS (Hoffman and Gelman, 2014)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aUjjeHOtxzw9"
      },
      "outputs": [],
      "source": [
        "def logistic_regression(features):\n",
        "  \"\"\"Bayesian logistic regression, which returns labels given features.\"\"\"\n",
        "  coeffs = ed.MultivariateNormalDiag(\n",
        "      loc=tf.zeros(features.shape[1]), name=\"coeffs\")\n",
        "  labels = ed.Bernoulli(\n",
        "      logits=tf.tensordot(features, coeffs, [[1], [0]]), name=\"labels\")\n",
        "  return labels\n",
        "\n",
        "features = tf.random_uniform([500, 55])\n",
        "true_coeffs = 5. * tf.random_normal([55])\n",
        "labels = tf.cast(tf.tensordot(features, true_coeffs, [[1], [0]]) \u003e 0,\n",
        "                 dtype=tf.int32)\n",
        "\n",
        "log_joint = ed.make_log_joint_fn(logistic_regression)\n",
        "def target_log_prob_fn(coeffs):\n",
        "  return log_joint(features=features, coeffs=coeffs, labels=labels)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "VaZpL6I2y--2"
      },
      "outputs": [],
      "source": [
        "if not tf.executing_eagerly():\n",
        "  raise ValueError(\"This code snippet requires eager execution.\")\n",
        "\n",
        "coeffs_samples = []\n",
        "target_log_prob = None\n",
        "grads_target_log_prob = None\n",
        "for step in range(500):\n",
        "  [\n",
        "      [coeffs],\n",
        "      target_log_prob,\n",
        "      grads_target_log_prob,\n",
        "  ] = kernel(target_log_prob_fn=target_log_prob_fn,\n",
        "             current_state=[coeffs],\n",
        "             step_size=[0.1],\n",
        "             current_target_log_prob=target_log_prob,\n",
        "             current_grads_target_log_prob=grads_target_log_prob)\n",
        "  coeffs_samples.append(coeffs)\n",
        "\n",
        "for coeffs_sample in coeffs_samples:\n",
        "  plt.plot(coeffs_sample.numpy())\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zhBRN_gxVUM6"
      },
      "source": [
        "See [`no_u_turn_sampler/logistic_regression.py`](https://github.com/google-research/google-research/tree/master/simple_probabilistic_programming/no_u_turn_sampler/logistic_regression.py) for the full example.\n",
        "\n",
        "## References\n",
        "\n",
        "1. Hoffman,  M. D. (2017).  Learning deep latent Gaussian models with Markov chain Monte Carlo. In _International Conference on Machine Learning_.\n",
        "2. Hoffman,  M. D. and Gelman,  A. (2014).  The No-U-turn sampler:  Adaptively setting path lengths in Hamiltonian Monte Carlo. _Journal of Machine Learning Research_, 15(1):1593–1623.\n",
        "3. Kusner, M. J., Paige, B., and Hernández-Lobato, J. M. (2017).  Grammar variational auto-encoder. In _International Conference on Machine Learning_.\n",
        "4. Papamakarios, G., Murray, I., and Pavlakou, T. (2017). Masked autoregressive flow for density estimation. In _Neural Information Processing Systems_.\n",
        "5. Parmar, N., Vaswani, A., Uszkoreit, J., Kaiser, Ł., Shazeer, N., Ku, A., and Tran, D. (2018). Image transformer. In _International Conference on Machine Learning_.\n",
        "6. Ranganath, R., Altosaar, J., Tran, D., and Blei, D. M. (2016). Operator variational inference. In _Neural Information Processing Systems_.\n",
        "7. Salimans, T., Kingma, D., and Welling, M. (2015). Markov chain Monte Carlo and variational inference: Bridging the gap. In _International Conference on Machine Learning_.\n",
        "8. Tran, D., Hoffman, M. D., Moore, D., Suter, C., Vasudevan S., Radul A., Johnson M., and Saurous R. A. (2018). Simple, Distributed, and Accelerated Probabilistic Programming. In _Neural Information Processing Systems_."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "4zzlF9Mt_DSG"
      ],
      "last_runtime": {
        "build_target": "",
        "kind": "local"
      },
      "name": "Companion.ipynb",
      "provenance": [
        {
          "file_id": "1ULybenz3Z-bnthUecd1wHUp7ugsZ0cAE",
          "timestamp": 1533574154382
        }
      ],
      "version": "0.3.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
